import librosa, os, math
import numpy as np
import matplotlib.pyplot as plt
from HyperParm import HyperParm
import random


class AudioParser:
    def __init__(
        self,
        sample_rate=None,
        fmin: int = None,
        fmax: int = None,
        n_fft: int = None,
        hop_length: int = None,
    ):
        self.audio: np.ndarray = None
        # hyperParam
        self.sample_rate = HyperParm.SAMPLE_RATE if sample_rate is None else sample_rate
        self.fmin = HyperParm.FMIN if fmin is None else fmin
        self.fmax = HyperParm.FMAX if fmax is None else fmax
        self.n_fft = HyperParm.N_FFT if n_fft is None else n_fft
        self.hop_length = HyperParm.HOP_LENGTH if hop_length is None else hop_length
        self.n_mels = HyperParm.N_MELS if hop_length is None else hop_length

    def gain_adjustment(self, target_lufs):

        # 计算当前响度
        current_lufs = 10 * np.log10(
            np.mean(librosa.perceptual_weighting(self.audio, self.sample_rate) ** 2)
        )

        # 计算增益 (dB)
        gain_dB = target_lufs - current_lufs
        gain = 10 ** (gain_dB / 20)  # dB 转线性增益，电压前面系数20

        self.audio *= gain

    def tile_audio(
        self, target_length: int, threshold: float = HyperParm.TILE_THRESHOLD
    ) -> bool:
        """循环填充autio，如果原长/目标长小于阈值，则失败；否则循环填充。返回是否成功"""
        length = len(self.audio)
        if length / target_length < threshold:
            return False

        self.audio = np.tile(self.audio, target_length // len(self.audio) + 1)[
            :target_length
        ]
        return True

    def split_audio(self, num: int):
        return np.array_split(self.audio, num)

    def load_audio(self, audio_path: str, mono=True):
        self.audio, _ = librosa.load(audio_path, sr=self.sample_rate, mono=mono)

    def get_mel(self, audio: np.ndarray = None):
        if audio is None:
            audio = self.audio
        stft_pic = librosa.stft(audio, n_fft=self.n_fft, hop_length=self.hop_length)
        power_spectrum = np.abs(stft_pic) ** 2  # 这里abs是复数取模

        mel_spectrogram = librosa.feature.melspectrogram(
            S=power_spectrum,
            sr=self.sample_rate,
            n_mels=self.n_mels,
            fmin=self.fmin,
            fmax=self.fmax,
        )

        # 转换为对数幅度
        log_mel_spectrogram = librosa.power_to_db(mel_spectrogram, ref=np.max).astype(
            HyperParm.MEL_DATA_TYPE
        )

        return log_mel_spectrogram

    def show_mel(self, mel: np.ndarray):
        plt.figure(figsize=(10, 6))
        librosa.display.specshow(
            mel,
            sr=self.sample_rate,
            y_axis="mel",
            fmin=self.fmin,
            fmax=self.fmax,
            x_axis="time",
            cmap="coolwarm",
        )
        plt.colorbar(format="%+2.0f dB")
        plt.title("Mel Spectrogram")
        plt.tight_layout()
        plt.show()

    def shuffle_audio(self, num: int):
        self.tile_audio(math.ceil(len(self.audio) / num) * num)
        self.audio.reshape((num, -1))
        np.random.shuffle(self.audio)
        self.audio.reshape(-1)

    def roll_audio(self, offset: int):
        self.audio = np.roll(self.audio, offset)

    def get_chroma(self, audio: np.ndarray = None):
        if audio is None:
            audio = self.audio
        return librosa.feature.chroma_stft(sr=self.sample_rate, y=audio)

    def get_spectral_contrast(self, audio: np.ndarray = None):
        if audio is None:
            audio = self.audio
        return librosa.feature.spectral_contrast(
            y=audio,
            sr=self.sample_rate,
            fmin=self.fmin,
            n_bands=HyperParm.N_BINDS,
        )


def test():

    root_path = os.getcwd()

    audioParser = AudioParser()
    audioParser.load_audio(
        "/home/tuchunxu/workspace/pr-project/pr-project/data/raw/blues/blues.00001.wav"
    )
    audioParser.get_spectral_contrast()
    # audioParser.gain_adjustment(0)
    # mel = audioParser.get_mel()
    # audioParser.show_mel(mel)

    # # 原始数组
    # arr = np.array([1, 2, 3])

    # # 目标长度
    # target_length = 10

    # # 使用 np.tile 填充
    # result = np.tile(arr, target_length // len(arr) + 1)[:target_length]

    # tmp = np.arange(0, 100)
    # # tmp = np.reshape(tmp, (10, 10))

    # print("循环填充结果:", tmp[:0])


if __name__ == "__main__":
    test()
